package cn.cloud.all.security.web.header.writers;

import cn.cloud.all.security.web.header.HeaderWriter;
import org.springframework.util.Assert;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

public class ReferrerPolicyHeaderWriter implements HeaderWriter {

    private static final String REFERRER_POLICY_HEADER = "Referrer-Policy";

    private ReferrerPolicy policy;

    /**
     * Creates a new instance. Default value: no-referrer.
     */
    public ReferrerPolicyHeaderWriter() {
        this(ReferrerPolicy.NO_REFERRER);
    }

    /**
     * Creates a new instance.
     *
     * @param policy a referrer policy
     * @throws IllegalArgumentException if policy is null
     */
    public ReferrerPolicyHeaderWriter(ReferrerPolicy policy) {
        setPolicy(policy);
    }

    /**
     * Sets the policy to be used in the response header.
     *
     * @param policy a referrer policy
     * @throws IllegalArgumentException if policy is null
     */
    public void setPolicy(ReferrerPolicy policy) {
        Assert.notNull(policy, "policy can not be null");
        this.policy = policy;
    }

    @Override
    public void writeHeaders(HttpServletRequest request, HttpServletResponse response) {
        if (!response.containsHeader(REFERRER_POLICY_HEADER)) {
            response.setHeader(REFERRER_POLICY_HEADER, this.policy.getPolicy());
        }
    }

    private enum ReferrerPolicy {

        NO_REFERRER("no-referrer"),
        NO_REFERRER_WHEN_DOWNGRADE("no-referrer-when-downgrade"),
        SAME_ORIGIN("same-origin"),
        ORIGIN("origin"),
        STRICT_ORIGIN("strict-origin"),
        ORIGIN_WHEN_CROSS_ORIGIN("origin-when-cross-origin"),
        STRICT_ORIGIN_WHEN_CROSS_ORIGIN("strict-origin-when-cross-origin"),
        UNSAFE_URL("unsafe-url");

        private static final Map<String, ReferrerPolicy> REFERRER_POLICIES;

        static {
            Map<String, ReferrerPolicy> referrerPolicies = new HashMap<>();
            for (ReferrerPolicy referrerPolicy : values()) {
                referrerPolicies.put(referrerPolicy.getPolicy(), referrerPolicy);
            }
            REFERRER_POLICIES = Collections.unmodifiableMap(referrerPolicies);
        }

        private String policy;

        ReferrerPolicy(String policy) {
            this.policy = policy;
        }

        public String getPolicy() {
            return this.policy;
        }

        public static ReferrerPolicy get(String referrerPolicy) {
            return REFERRER_POLICIES.get(referrerPolicy);
        }
    }
}
